#!/usr/bin/env python
# coding=utf-8
# __author__ = 'Yunchao Ling'

from pymongo import MongoClient
import os
import datetime
import time
import re
import sys

HOST = 'localhost'
PORT = 27017
DB_NAME = 'pubmed'
COLLECTION_NAME = 'metalist'


def GetConfig():
    if os.path.exists("../config.ini"):
        import ConfigParser
        global HOST
        global PORT
        global DB_NAME
        global COLLECTION_NAME
        config = ConfigParser.ConfigParser()
        config.read("../config.ini")
        HOST = config.get("mongodb", "host")
        PORT = int(config.get("mongodb", "port"))
        DB_NAME = config.get("mongodb", "db")
        COLLECTION_NAME = config.get("mongodb", "collection")


def GetPMID(line):
    pattern = re.compile("<.*>(.*)<.*>", re.I)
    m = re.search(pattern, line)
    if m:
        return m.group(1)
    else:
        return m


# Use regex to analyze PMID lines
def Extract_PubMed_IDs_From_XML(filename):
    infile = open(filename, "r")
    result = []
    for line in infile:
        line = line.rstrip()
        if line.startswith("        <PMID"):
            pmid_string = GetPMID(line)
            if pmid_string.isdigit() and pmid_string != None:
                result.append(pmid_string)
    infile.close()
    return result


if __name__ == '__main__':
    os.chdir(sys.argv[1])

    print("Loading date, given date is ... "),
    time.sleep(1)
    dateString = sys.argv[1].split("_")[-1]
    dateObject = datetime.datetime.strptime(dateString, "%Y%m%d")
    print(dateString + ".")
    time.sleep(1)

    print("Loading config file ... "),
    time.sleep(1)
    GetConfig()
    print("complete.")
    time.sleep(1)

    print("Initiating mongodb connection ... "),
    time.sleep(1)
    connection = MongoClient(HOST, PORT)
    db = connection[DB_NAME]
    collection = db[COLLECTION_NAME]
    print("complete.")
    time.sleep(1)

    print("Loading PubMed xml files ... "),
    time.sleep(1)
    filelist = os.listdir(".")
    print("complete.")
    time.sleep(1)

    print("Start anlyzing xml files ... ")
    time.sleep(1)
    for filename in filelist:
        if filename.endswith(".xml"):
            print("Start analyzing " + filename + " ... "),
            filepath = os.path.dirname(os.path.abspath(filename)).split("/")[-1] + "/" + filename
            ids = Extract_PubMed_IDs_From_XML(filename)
            for id in ids:
                pmid = int(id)
                collection.update_one({"pmid": pmid},
                                      {"$set": {"xml": filepath, "time_in": dateObject, "time_last": dateObject}},
                                      upsert=True)
            print("complete. Find " + str(len(ids)) + " PMIDs.")
    connection.close()
    print("All xml files analyzed.")
